﻿module Coloring (color,check_colors,get_colors) where

import List (sortBy, find, delete, elem)
import Utils (add_in_set_mode)
import Data.Graph.Inductive.Graph (Node, outdeg, labNodes, LNode, edges, mkGraph, lab, nodes, neighbors)
import Data.Graph.Inductive.Tree (Gr, UGr)
--import GraphDef
import Data.Graph.Inductive.Example (labUEdges)
import System.CPUTime (getCPUTime)

--Gets colors if the graph		
get_colors:: Gr Int () -> [Int]
get_colors = foldl (flip $ (flip add_in_set_mode) . snd) [] . labNodes

--Coloring of the graph
color :: UGr -> (Gr Int (),String,Int)
color gr = color_inner split_by_colors iterLim True ""  where
	color_inner :: [(Int,[Node])] -> Int -> Bool -> String -> (Gr Int (),String,Int)
	color_inner groups 0 _ trace = (apply_colors groups, trace, length groups)
	color_inner groups iterNum bySize trass
		| (length groups) == (length sortedGroups) = 
			color_inner sortedGroups (iterNum-1) (not bySize) newTrass
		| otherwise = 
			color_inner sortedGroups iterLim (not bySize) newTrass
		where
		sortedGroups = peform_step $ (if bySize then sort_by_size else sort_by_deg) groups
		newTrass = concat [trass, groups_to_string groups, " iter: ", show iterNum, "\n"]
		
	grNodes = nodes gr
	grEdges = edges gr
	grLEdges = labUEdges grEdges
	nbs_f = neighbors gr

	--Sorts groups of nodes by their size
	sort_by_size :: [(Int,[Node])] -> [(Int,[Node])]
	sort_by_size = sortBy $ (flip $ compare . length . snd) . length . snd

	--Sorts groups of nodes by their total degree
	sort_by_deg :: [(Int,[Node])] -> [(Int,[Node])]
	sort_by_deg = sortBy $ (flip $ compare . sum_deg . snd) . sum_deg . snd where
		sum_deg = (`foldl` 0) $ flip $ (+) . (outdeg gr)

	--Splits nodes of the graph by their color into groups
	split_by_colors :: [(Int,[Node])]
	split_by_colors = color_nodes 0 grNodes [] where

		color_nodes :: Int -> [Int] -> [(Int,[Node])] -> [(Int,[Node])]
		color_nodes _ [] groups = groups		
		color_nodes num lst@(node:nodes) groups
			| foldl (\ flag node -> flag || eig node) False nbs = color_nodes (num+1) lst groups
			| otherwise = color_nodes 0 nodes (insert_node node num groups)
			where
			nbs = nbs_f node	
			eig = exist_in_groups groups num False

		insert_node :: Node -> Int -> [(Int,[Node])] -> [(Int,[Node])]
		insert_node node label groups = 
			case (`find` groups) $ (== label) . fst of
			Nothing -> (label,[node]):groups
			Just gr@(lab,group) -> (lab,node:group):(delete gr groups)
			
		exist_in_groups :: [(Int,[Node])] -> Int -> Bool -> Node -> Bool
		-- exists_in_groups _ _ True _ = True
		exist_in_groups [] _ flag _ = flag
		exist_in_groups ((color,nodes):leftGroups) label flag node = exist_in_groups leftGroups label (flag||((label == color) && (elem node nodes))) node

	groups_to_string :: [(Int,[Node])] -> String
	groups_to_string = foldl (\str (label,nodes) -> concat [str, "Color ", show label, ": ", node_list_to_string nodes, "\n"]) "" where
		node_list_to_string = (`foldl` "") $ flip $ (++) . (++ " "). show

	--One step of algorithm
	peform_step :: [(Int,[Node])] -> [(Int,[Node])]
	peform_step [] = error "Empty graph"
	peform_step (lgroup:groups) = peform_step_inner lgroup [] groups [] where
		peform_step_inner :: (Int,[Node]) -> [Node] -> [(Int,[Node])] -> [(Int,[Node])] -> [(Int,[Node])]
		peform_step_inner (lab,[]) doneNodes [] = ((lab, doneNodes):)
		peform_step_inner (lab,[]) doneNodes (lgroup:leftGroups) = 
			(peform_step_inner lgroup [] leftGroups) . (if (doneNodes /= []) then ((lab, doneNodes):) else id)
		peform_step_inner (lab,(node:leftNodes)) doneNodes leftGroups =
			case insert_node node leftGroups [] of
			Nothing -> peform_step_inner (lab,leftNodes) (node: doneNodes) leftGroups
			Just grs -> peform_step_inner (lab,leftNodes) doneNodes grs
	
		--Insertion node in groups
		insert_node :: Node -> [(Int,[Node])] -> [(Int,[Node])] -> Maybe [(Int,[Node])]
		insert_node _ [] = \_ -> Nothing
		insert_node node (lgroup@(lab,group):leftGroups)
			| need_insert_node group = Just . (++ (lab,node:group):leftGroups) . reverse
			| otherwise  = (insert_node node leftGroups) . (lgroup:)
			where
			nbs = nbs_f node
			
			need_insert_node :: [Int] -> Bool
			need_insert_node = foldr ((&&) . not . (flip elem) nbs) True

	apply_colors :: [(Int,[Node])] -> Gr Int ()
	apply_colors = 
		(flip mkGraph) grLEdges . foldl (\nodeList (label,group) -> (zip group $ repeat label) ++ nodeList) []

	--number of unproductive iterations
	iterLim = 10
			
--Checking coloring of the graph
check_colors :: Gr Int() -> Bool
check_colors graph = foldl (\valid (start,end) -> valid && (lab graph start) /= (lab graph end)) True (edges graph)
